// SPDX-FileCopyrightText: 2025 UnionTech Software Technology Co., Ltd.
// SPDX-License-Identifier: MIT

package metrics

import (
	"context"
	"fmt"
	"math"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

const (
	chronysubsystem      = "chrony"
	chronyCommandTimeout = 5 * time.Second
)

type ChronyCollector struct {
	config  ChronyConfig
	metrics map[string]*MetricDesc
}

func NewChronyCollector(cfg *ChronyConfig) *ChronyCollector {
	return &ChronyCollector{
		config: *cfg,
		metrics: map[string]*MetricDesc{
			"trackingLI": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_li"),
					"Chrony tracking LI",
					nil,
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingClockSource": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_clock_source"),
					"Chrony tracking clock source",
					[]string{"source"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingCorrection": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_correction"),
					"Chrony tracking correction",
					nil,
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingStratum": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_stratum"),
					"Chrony tracking stratum",
					nil,
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingPrecision": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_precision"),
					"Chrony tracking precision",
					nil,
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingRootDelay": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_root_delay"),
					"Chrony tracking root delay",
					nil,
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingRootDisp": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_root_disp"),
					"Chrony tracking root dispersion",
					nil,
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingRefID": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_ref_id"),
					"Chrony tracking ref id",
					[]string{"ref_id"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingRefTime": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_ref_time"),
					"Chrony tracking ref time",
					nil,
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingOffset": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_offset"),
					"Chrony tracking offset",
					nil,
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"trackingFrequency": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "tracking_frequency"),
					"Chrony tracking frequency",
					nil,
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			// sources metrics
			"sourcesStratum": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sources_stratum"),
					"Chrony source stratum",
					[]string{"source_address", "mode", "state"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcesPoll": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sources_poll"),
					"Chrony source poll interval",
					[]string{"source_address", "mode", "state"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcesReach": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sources_reach"),
					"Chrony source reach register",
					[]string{"source_address", "mode", "state"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcesLastRx": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sources_last_rx"),
					"Chrony source last receive",
					[]string{"source_address", "mode", "state"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcesLastSample": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sources_last_sample"),
					"Chrony source last sample offset",
					[]string{"source_address", "mode", "state"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			// sourcestats metrics
			"sourcestatsNpSamples": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sourcestats_np_samples"),
					"Chrony source stats number of samples",
					[]string{"source_address"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcestatsNpRuns": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sourcestats_np_runs"),
					"Chrony source stats number of runs",
					[]string{"source_address"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcestatsSpanSeconds": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sourcestats_span_seconds"),
					"Chrony source stats span seconds",
					[]string{"source_address"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcestatsFrequency": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sourcestats_frequency"),
					"Chrony source stats frequency",
					[]string{"source_address"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcestatsFreqSkew": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sourcestats_freq_skew"),
					"Chrony source stats frequency skew",
					[]string{"source_address"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcestatsOffset": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sourcestats_offset"),
					"Chrony source stats offset",
					[]string{"source_address"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
			"sourcestatsStdDev": {
				Desc: prometheus.NewDesc(
					prometheus.BuildFQName(namespace, chronysubsystem, "sourcestats_std_dev"),
					"Chrony source stats standard deviation",
					[]string{"source_address"},
					nil,
				),
				Type: prometheus.GaugeValue,
			},
		},
	}
}

func (c *ChronyCollector) Update(ch chan<- prometheus.Metric) error {
	metrics, err := c.collectChronyMetrics()
	if err != nil {
		return err
	}

	// 发送tracking指标
	c.sendTrackingMetrics(metrics.Tracking, ch)

	// 发送sources指标
	c.sendSourcesMetrics(metrics.Sources, ch)

	// 发送sourcestats指标
	c.sendSourceStatsMetrics(metrics.SourceStats, ch)

	return nil
}

// 主收集逻辑
func (c *ChronyCollector) collectChronyMetrics() (*ChronyMetrics, error) {
	if !c.isChronyAvailable() {
		return nil, ErrNoData
	}

	metrics := &ChronyMetrics{}

	// 收集tracking数据
	if err := c.collectTracking(metrics); err != nil {
		return nil, err
	}

	// 收集sources数据
	if err := c.collectSources(metrics); err != nil {
		logrus.Warnf("Failed to collect chrony sources: %v", err)
	}

	// 收集sourcestats数据
	if err := c.collectSourceStats(metrics); err != nil {
		logrus.Warnf("Failed to collect chrony sourcestats: %v", err)
	}

	return metrics, nil
}

type ChronyMetrics struct {
	Tracking    TrackingMetrics
	Sources     []SourceMetrics
	SourceStats []SourceStatsMetrics
}

type TrackingMetrics struct {
	LI          uint8
	ClockSource string
	Correction  float64
	Stratum     int
	Precision   int
	RootDelay   float64
	RootDisp    float64
	RefID       string
	RefTime     string
	Offset      float64
	Frequency   float64
}

type SourceMetrics struct {
	SourceAddress string
	Mode          string
	State         string
	Stratum       int
	Poll          int
	Reach         int
	LastRx        int
	LastSample    float64
}

type SourceStatsMetrics struct {
	SourceAddress string
	NpSamples     int
	NpRuns        int
	SpanSeconds   int
	Frequency     float64
	FreqSkew      float64
	Offset        float64
	StdDev        float64
}

func (c *ChronyCollector) isChronyAvailable() bool {
	_, err := exec.LookPath("chronyc")
	return err == nil
}

func (c *ChronyCollector) execChronyc(command string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), chronyCommandTimeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, "chronyc", "-c", command)
	output, err := cmd.CombinedOutput()
	if err != nil {
		if ctx.Err() == context.DeadlineExceeded {
			return "", fmt.Errorf("chronyc command timeout after %s", chronyCommandTimeout)
		}
		return "", err
	}
	return string(output), nil
}

func (c *ChronyCollector) collectTracking(metrics *ChronyMetrics) error {
	output, err := c.execChronyc("tracking")
	if err != nil {
		return err
	}

	fields := strings.Split(output, ",")
	if len(fields) < 14 {
		return ErrNoData
	}

	metrics.Tracking = TrackingMetrics{
		RefID:       fields[0],
		Stratum:     parseInt(fields[2]),
		RefTime:     fields[3],
		Correction:  parseFloat(fields[4]),
		LI:          parseLeapStatus(fields[13]),
		Offset:      parseFloat(fields[5]),
		Frequency:   parseFloat(fields[7]),
		RootDelay:   parseFloat(fields[10]),
		RootDisp:    parseFloat(fields[11]),
		ClockSource: "ntp",
		Precision:   -int(math.Log2(parseFloat(fields[4]))),
	}
	return nil
}

func (c *ChronyCollector) sendTrackingMetrics(t TrackingMetrics, ch chan<- prometheus.Metric) {
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingLI"].Desc,
		prometheus.GaugeValue,
		float64(t.LI),
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingClockSource"].Desc,
		prometheus.GaugeValue,
		1,
		t.ClockSource,
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingCorrection"].Desc,
		prometheus.GaugeValue,
		t.Correction,
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingStratum"].Desc,
		prometheus.GaugeValue,
		float64(t.Stratum),
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingPrecision"].Desc,
		prometheus.GaugeValue,
		float64(t.Precision),
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingRootDelay"].Desc,
		prometheus.GaugeValue,
		t.RootDelay,
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingRootDisp"].Desc,
		prometheus.GaugeValue,
		t.RootDisp,
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingRefID"].Desc,
		prometheus.GaugeValue,
		1,
		t.RefID,
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingRefTime"].Desc,
		prometheus.GaugeValue,
		parseFloat(t.RefTime),
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingFrequency"].Desc,
		prometheus.GaugeValue,
		t.Frequency,
	)
	ch <- prometheus.MustNewConstMetric(
		c.metrics["trackingOffset"].Desc,
		prometheus.GaugeValue,
		t.Offset,
	)
}

func parseFloat(s string) float64 {
	v, _ := strconv.ParseFloat(strings.TrimSpace(s), 64)
	return v
}

func parseInt(s string) int {
	v, _ := strconv.Atoi(strings.TrimSpace(s))
	return v
}

func parseLeapStatus(s string) uint8 {
	switch strings.ToLower(strings.TrimSpace(s)) {
	case "normal":
		return 0
	case "leapsecond":
		return 1
	case "notsynchronized":
		return 2
	default:
		logrus.Warnf("Unknown leap status: %s", s)
		return 3
	}
}

func (c *ChronyCollector) collectSources(metrics *ChronyMetrics) error {
	output, err := c.execChronyc("sources")
	if err != nil {
		return err
	}

	lines := strings.Split(output, "\n")
	metrics.Sources = make([]SourceMetrics, 0)

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		fields := strings.Split(line, ",")
		if len(fields) < 8 {
			continue
		}

		source := SourceMetrics{
			SourceAddress: fields[0],
			Mode:          fields[1],
			State:         fields[2],
			Stratum:       parseInt(fields[3]),
			Poll:          parseInt(fields[4]),
			Reach:         parseInt(fields[5]),
			LastRx:        parseInt(fields[6]),
			LastSample:    parseFloat(fields[7]),
		}

		metrics.Sources = append(metrics.Sources, source)
	}

	return nil
}

func (c *ChronyCollector) sendSourcesMetrics(sources []SourceMetrics, ch chan<- prometheus.Metric) {
	for _, source := range sources {
		labels := []string{source.SourceAddress, source.Mode, source.State}

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcesStratum"].Desc,
			prometheus.GaugeValue,
			float64(source.Stratum),
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcesPoll"].Desc,
			prometheus.GaugeValue,
			float64(source.Poll),
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcesReach"].Desc,
			prometheus.GaugeValue,
			float64(source.Reach),
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcesLastRx"].Desc,
			prometheus.GaugeValue,
			float64(source.LastRx),
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcesLastSample"].Desc,
			prometheus.GaugeValue,
			source.LastSample,
			labels...,
		)
	}
}

func (c *ChronyCollector) collectSourceStats(metrics *ChronyMetrics) error {
	output, err := c.execChronyc("sourcestats")
	if err != nil {
		return err
	}

	lines := strings.Split(output, "\n")
	metrics.SourceStats = make([]SourceStatsMetrics, 0)

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		fields := strings.Split(line, ",")
		if len(fields) < 8 {
			continue
		}

		sourceStat := SourceStatsMetrics{
			SourceAddress: fields[0],
			NpSamples:     parseInt(fields[1]),
			NpRuns:        parseInt(fields[2]),
			SpanSeconds:   parseInt(fields[3]),
			Frequency:     parseFloat(fields[4]),
			FreqSkew:      parseFloat(fields[5]),
			Offset:        parseFloat(fields[6]),
			StdDev:        parseFloat(fields[7]),
		}

		metrics.SourceStats = append(metrics.SourceStats, sourceStat)
	}

	return nil
}

func (c *ChronyCollector) sendSourceStatsMetrics(sourceStats []SourceStatsMetrics, ch chan<- prometheus.Metric) {
	for _, stat := range sourceStats {
		labels := []string{stat.SourceAddress}

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcestatsNpSamples"].Desc,
			prometheus.GaugeValue,
			float64(stat.NpSamples),
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcestatsNpRuns"].Desc,
			prometheus.GaugeValue,
			float64(stat.NpRuns),
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcestatsSpanSeconds"].Desc,
			prometheus.GaugeValue,
			float64(stat.SpanSeconds),
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcestatsFrequency"].Desc,
			prometheus.GaugeValue,
			stat.Frequency,
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcestatsFreqSkew"].Desc,
			prometheus.GaugeValue,
			stat.FreqSkew,
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcestatsOffset"].Desc,
			prometheus.GaugeValue,
			stat.Offset,
			labels...,
		)

		ch <- prometheus.MustNewConstMetric(
			c.metrics["sourcestatsStdDev"].Desc,
			prometheus.GaugeValue,
			stat.StdDev,
			labels...,
		)
	}
}
